import math
import time
from modis_utils import *


def merge_crop_info(input_path, mark_path, output_path, fill):
    #input data
    inds_input = gdal.Open(input_path)
    geotransform_g = inds_input.GetGeoTransform()
    nxsize_input = inds_input.RasterXSize
    nysize_input = inds_input.RasterYSize
    im_data_input = inds_input.ReadAsArray(0, 0, nxsize_input, nysize_input)
    #mark data
    inds_mark = gdal.Open(mark_path)
    geotransform_mark = inds_mark.GetGeoTransform()
    nxsize_mark = inds_mark.RasterXSize
    nysize_mark = inds_mark.RasterYSize
    im_data_mark = inds_mark.ReadAsArray(0, 0, nxsize_mark, nysize_mark)
    in_band_mark = inds_mark.GetRasterBand(1)

    merge_data = np.zeros((nysize_input, nxsize_input))
    startXpoint = math.floor((geotransform_g[0] - geotransform_mark[0]) / geotransform_g[1])
    startYpoint = math.floor((geotransform_g[3] - geotransform_mark[3]) / geotransform_g[5])
    for i in range(0, nysize_input):
        for j in range(0, nxsize_input):
            ii = i+startYpoint
            jj = j+startXpoint
            if im_data_input[i, j] == 11:
                merge_data[i, j] = 11
            elif im_data_input[i, j] in land_list and im_data_mark[ii, jj] in irrigated_list:
                merge_data[i, j] = 14
            elif im_data_input[i, j] in land_list and im_data_mark[ii, jj] in rainfed_list:
                merge_data[i, j] = 15
            elif im_data_mark[ii, jj] == 9:
                merge_data[i, j] = 9
            elif im_data_mark[ii, jj] == 0:
                merge_data[i, j] = 0
    #create output data
    driver = gdal.GetDriverByName('GTiff')
    out_ds = driver.Create(output_path, nxsize_input, nysize_input, 1, in_band_mark.DataType)
    out_ds.SetGeoTransform(geotransform_g)
    out_ds.SetProjection(inds_input.GetProjection())
    out_ds.GetRasterBand(1).SetNoDataValue(fill)
    out_ds.GetRasterBand(1).Fill(fill)
    out_ds.GetRasterBand(1).SetRasterColorTable(set_color_table())
    t_band = out_ds.GetRasterBand(1)
    t_band.WriteArray(merge_data)
    del driver, inds_input, inds_mark


def tiff_warp(tiff_list, output_path):
    for in_fns in tiff_list:
        print("start warp {}".format(in_fns))
        (filename, extension) = os.path.splitext(in_fns)
        output_filename = output_path +'/{}_conver.tif'.format(filename.split('\\')[-1])
        in_ds = gdal.Open(in_fns)
        gdal.Warp(output_filename, in_ds, dstSRS='EPSG:4326', xRes=0.5, yRes=0.5, resampleAlg=gdal.GRA_NearestNeighbour)


def merge_tiff(tiff_list, output_dir, fill):
    min_x, max_y, max_x, min_y = get_extent(tiff_list[0])
    for in_fn in tiff_list[1:]:
        minx, maxy, maxx, miny = get_extent(in_fn)
        min_x = min(min_x, minx)
        min_y = min(min_y, miny)
        max_x = max(max_x, maxx)
        max_y = max(max_y, maxy)
    in_ds = gdal.Open(tiff_list[0])
    geotrans = list(in_ds.GetGeoTransform())
    geotrans = [min_x, geotrans[1], 0, max_y, 0, geotrans[5]]
    xsize = int((max_x - min_x) / geotrans[1] + 0.5)
    ysize = int((min_y - max_y) / geotrans[5] + 0.5)
    driver = gdal.GetDriverByName('GTiff')
    out_ds = driver.Create(output_dir, xsize, ysize, 1, gdal.GDT_Float32)
    out_ds.SetGeoTransform(geotrans)
    out_ds.SetProjection(in_ds.GetProjection())
    out_ds.GetRasterBand(1).SetNoDataValue(fill)
    out_ds.GetRasterBand(1).Fill(fill)
    for in_fns in tiff_list:
        in_ds = gdal.Open(in_fns)
        copy_into(out_ds, in_ds, 1, 1, fill)


if __name__ == '__main__':
    fill = 0
    year = 2010
    input_path = 'F:/16satellite/data/gongpeng/new/globalLC_250m_mod_2010'
    output_path = 'F:/16satellite/data/gongpeng/crop_info_dir'
    mark_path = 'F:/16satellite/data/download/GFSAD1KCD.tif'
    tiff_list = find_layers_files(input_path, '.tif')
    output_path_temp = output_path + '/temp'
    start_time = time.time()
    if not os.path.exists(output_path_temp):
        os.makedirs(output_path_temp)
    tiff_warp(tiff_list, output_path_temp)
    output_tif_path = output_path + '/{}_merge.tif'.format(year)
    print(output_tif_path)
    conver_tiff_list = find_layers_files(output_path_temp, '.tif')
    if os.path.exists(output_tif_path):
        os.remove(output_tif_path)
        print('{} tiff file is exist and remove'.format(output_tif_path))
    merge_tiff(conver_tiff_list, output_tif_path, fill)
    merge_info_path = output_path + '/{}_crop_info.tif'.format(year)
    if os.path.exists(merge_info_path):
        os.remove(merge_info_path)
        print('{} tiff file is exist and remove'.format(merge_info_path))
    merge_crop_info(output_tif_path, mark_path, merge_info_path, fill)
    if os.path.exists(output_path_temp):
        del_file(output_path_temp)
        os.rmdir(output_path_temp)
    print('Complete one year crop information extraction, cost time is: {} seconds'.format(
        time.time() - start_time))